#pragma once

#include <iostream>
#include <concepts>
#include <coroutine>
#include <memory>
#include <vector>

// 模仿UE5Coro的框架

namespace UE5Coro {

class Promise;
class LatentPromise;

// Awaiter
class LatentAwaiter {
    // 判断继续的数据与方法
public:
    bool ShouldResume() { return Resume(State, false); };
    // template<std::derived_from<Promise> P>
    // void Suspend(P&);
    // 这里写模板的话，没法索引到LatentPromise，会索引到TCoroutinePromise，导致编译错误
    void Suspend(LatentPromise& Promise) noexcept;
protected:
    void* State;
    bool (*Resume)(void* State, bool bCleanup);
    // 原生Cpp Coroutine的接口
public:
    [[nodiscard]] bool await_ready() { return ShouldResume(); }
    
    template<std::derived_from<Promise> P>
    void await_suspend(std::coroutine_handle<P> Handle) {
        Suspend(Handle.promise());
    }
    void await_resume() noexcept {}
    
    // 构造函数
    LatentAwaiter(void* state, bool (*resume)(void*, bool)) : State(state), Resume(resume) {}
};
class FInitialSuspendAwaiter {
public:
    enum class EAction {
        Destory,
        Resume,
    };
    EAction Action;

    [[nodiscard]] bool await_ready() noexcept { return false; }
    template<std::derived_from<Promise> P>
    void await_suspend(std::coroutine_handle<P> Handle) noexcept {
        if (Action == EAction::Destory) {
            Handle.destroy();
        } else if (Action == EAction::Resume) {
            Handle.resume();
        }
    }
    void await_resume() noexcept {}
};
class FFinalSuspendAwaiter {
public:
    bool bDestory;

    [[nodiscard]] bool await_ready() noexcept { return false; }
    template<std::derived_from<Promise> P>
    void await_suspend(std::coroutine_handle<P> Handle) noexcept { 
        if (bDestory) {
            Handle.destroy();
        }
    }
    void await_resume() noexcept {}
};

// Simulate LatentAction
class LatentAction {
    public:
        bool bIsDone = false;
    
        LatentAction(const LatentAwaiter& awaiter, LatentPromise* promise) : CurrentAwaiter(awaiter), CurrentPromise(promise) {}
        ~LatentAction() {
            CurrentPromise = nullptr;
        }
        void UpdateOperation() {
            if (bIsDone) {
                return;
            }
            if (CurrentPromise == nullptr) {
                return;
            }
            if (CurrentAwaiter.ShouldResume()) {
                std::coroutine_handle<LatentPromise> Handle = 
                    std::coroutine_handle<LatentPromise>::from_promise(*CurrentPromise);
                Handle.resume();
                bIsDone = true;
            }
        }
    private:
        LatentAwaiter CurrentAwaiter;
        LatentPromise* CurrentPromise = nullptr;
};
class LatentActionManager {
private:
    std::vector<LatentAction*> LatentActions;
public:
    LatentActionManager() {}
    ~LatentActionManager() {
        for (LatentAction* Action : LatentActions) {
            delete Action;
        }
    }
    void AddLatentAction(LatentAction* Action) {
        LatentActions.push_back(Action);
    }
    void RemoveLatentAction(LatentAction* Action) {
        LatentActions.erase(std::remove(LatentActions.begin(), LatentActions.end(), Action), LatentActions.end());
    }
    void UpdateOperations() {
        for (LatentAction* Action : LatentActions) {
            Action->UpdateOperation();
        }
    }
    static LatentActionManager* Get() {
        static LatentActionManager Instance;
        return &Instance;
    }
};

// Coroutine - 前向声明
template<typename T = void>
class TCoroutine;
template<>
class TCoroutine<> {
public:
    void Resume() {}

    LatentAction* CurrentLatentAction;
    ~TCoroutine<>() {
        if (CurrentLatentAction != nullptr) {
            delete CurrentLatentAction;
            CurrentLatentAction = nullptr;
        }
    }
};
template<typename T>
class TCoroutine: public TCoroutine<> {
    using ResultType = T;
};
struct VoidCoroutine final : TCoroutine<> {
    std::coroutine_handle<LatentPromise> Handle;
    VoidCoroutine(std::coroutine_handle<LatentPromise> h) : Handle(h) {
        CurrentLatentAction = nullptr;
    }
};

// Promise
class Promise {
};
class LatentPromise : public Promise {
    // 标准Cpp Coroutine接口
public:
    // 构造函数
    LatentPromise() = default;
    LatentPromise(const LatentPromise&) = default;
    LatentPromise& operator=(const LatentPromise&) = default;
    
    FInitialSuspendAwaiter initial_suspend();           // 返回值需要实现Await接口
    FFinalSuspendAwaiter final_suspend() noexcept;      // 返回值需要实现Await接口
};
template<typename T, typename Base>
class TCoroutinePromise: public Base {
public:
    TCoroutine<T> get_return_object() {
        return TCoroutine<T>{};
    }
    void return_value(T) {
        // 存储返回值
    }
    void unhandled_exception() {
        // 处理异常
    }
};

template<typename Base>
class TCoroutinePromise<void, Base>: public Base {
public:
    VoidCoroutine get_return_object() {
        return VoidCoroutine{std::coroutine_handle<LatentPromise>::from_promise(*this)};
    }
    void return_void() {
        // void协程不需要返回值
    }
    void unhandled_exception() {
        // 处理异常
    }
};

void Test();
}

// 根据Cpp的规则，必须在std的namespace进行模板特化
template<typename T, typename... Args>
struct std::coroutine_traits<UE5Coro::TCoroutine<T>, Args...> {
    using promise_type = UE5Coro::TCoroutinePromise<T, UE5Coro::LatentPromise>;
};
template<typename... Args>
struct std::coroutine_traits<UE5Coro::VoidCoroutine, Args...> {
    using promise_type = typename coroutine_traits<UE5Coro::TCoroutine<>, Args...>::promise_type;
};
